from numpy import linspace, cos, sin, meshgrid, array, empty, pi, mgrid, tile
from tvtk.api import tvtk
from mayavi import mlab
from wx import Yield
from pdb import set_trace
"""
Module to create cylindrical static and animated plots from afd simulations.
The plots are created using mayavi module with StructuredGrid data TVTK 
datasets. 
Created: Brent Maranzano
Last Modified: June 2, 2014
"""

def _create_sg(r, z, nq):
    """
    Create a structured grid from the radial and axial coordinates.
    r : [nr x nz] numpy array of the radial coordinates
    z : [nr x nz] numpy array of the axial coordinates
    nq : integer number of points in the azimuthal coordinate
    return tvtk StructuredGrid
    """
    r1D = r[:, 0]
    z1D = z[0, :]
    q1D = linspace(0, 2*pi, nq)

    dims = (r1D.size, q1D.size, z1D.size) 
    pts = empty((dims[0], dims[1], dims[2], 3), dtype=float)

    q, r = meshgrid(q1D, r1D)

    x = r * cos(q)
    y = r * sin(q)

    for i, zp in enumerate(z[0,:]):
        pts[..., i, 0] = x
        pts[..., i, 1] = y
        pts[..., i, 2] = zp 

    pts = pts.transpose(2, 1, 0, 3).copy()
    pts.shape = pts.size / 3, 3

    sg = tvtk.StructuredGrid(dimensions=dims, points=pts)
    return sg

def set_sg_data(sg, data, name):
    """
    Set the data values from a simulation (temperature, liquid volume
    fraction) to the structured grid.
    sg : tvtk StructuredGrid for the data to be added
    data : [nr x nz] numpy array of the data values
    name : string name of the dataset
    return tvtk StructuredGrid with data set to data
    """
    dims = sg.dimensions
    data3D = empty((dims[0], dims[1], dims[2]), dtype=float)

    data3D = tile(data[:, None, :], (1, dims[1], 1))

    data3D = data3D.T.copy().ravel()

    sg.point_data.scalars = data3D
    sg.point_data.scalars.name = name

    return sg

def sg_plot(r ,z, nq, data, name):
    """
    Make a plot of the vtk structured grid data
    r : [nr x nz] numpy array of radial coordinates
    z : [nr x nz] numpy array of axial coordinates
    nq : integer number of azimuthal coordinates
    data : [nr x nz] numpy array of data values
    name : string name of data
    return mayavi scene
    """
    if len(data.shape) != 2:
        raise ValueError('data should be a [nr x nz] numpy array')

    f = mlab.figure()

    sg = _create_sg(r, z, nq)
    sg = set_sg_data(sg, data, name)
    d = mlab.pipeline.add_dataset(sg)

    surf = mlab.pipeline.surface(d, opacity=0.6)

    delta = data.max() - data.min()
    contours_ = [data.min() + frac*delta for frac in [0.1, 0.3, 0.5, 0.7, 0.9]]
    iso_surface = mlab.pipeline.iso_surface(d, opacity=0.2)#, contours=contours_)

    f.scene.background = (1.0, 1.0, 1.0)
    mlab.view(azimuth=0, elevation=60)

    mlab.show()

    return f

def animate(r ,z, nq, data, name):
    """
    Function to create time varying 3D graphics to depict
    the AFD simulation results.
    r : [nr x nz] numpy array of the radial coordinates
    z : [nr x nz] numpy array of the axial coordinates
    nq : integer number of points in the azimuthal coordinate
    data : [nt x nr x nz] numpy array of dta
    name : string name of data
    """
    f = sg_plot(r, z, nq, data[0], name)
    d = f.children[0]
    sg = d.data
    for dat in data:
        sg = set_sg_data(sg, dat, name)
        d.update()
        Yield()

@mlab.animate(ui=True, delay=100)
def ui_animate(r, z, nq, data, name):
    """
    Identical to the animate function except decorated to give delay control
    and a user interface.
    """
    f = sg_plot(r, z, nq, data[0], name)
    d = f.children[0]
    sg = d.data
    for dat in data:
        sg = set_sg_data(sg, dat, name)
        d.update()
        yield 

def bit_map(r ,z, nq, data, name):
    """
    Generator that yields a bitmap array of the mayavi plot.
    Same syntax as function animate, but yields the array for
    each slice of data.
    r : [nr x nz] numpy array of the radial coordinates
    z : [nr x nz] numpy array of the axial coordinates
    nq : integer number of points in the azimuthal coordinate
    data : [nt x nr x nz] numpy array of dta
    name : string name of data
    """
    f = sg_plot(r, z, nq, data[0], name)
    d = f.children[0]
    sg = d.data
    for dat in data:
        sg = set_sg_data(sg, dat, name)
        d.update()
        yield mlab.screenshot(figure=f)

def test():
    """
    Function to test the plotting program. Make sure that the
    wx backend has been set. For certainty execute
    :%gui wx
    at the python prompt.
    The results should show a cylinder with a parabolic isocontours.
    Usage:
    cylinder.test()
    """
    r, z = mgrid[2.63e-5:.05:20j, 0:.04:30j]
    data = array([10*i*r**2 + i*z for i in range(30)])
    ui_animate(r, z, 25, data, "test")

if __name__=="__main__":
    pass
